import torch
import torch.nn as nn
import os
from model import StrawberryModel


def create_sample_weights():
    """创建与模型结构完全匹配的权重文件"""
    # 创建模型实例
    model = StrawberryModel()

    # 初始化权重
    for m in model.modules():
        if isinstance(m, nn.Conv2d):
            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        elif isinstance(m, nn.BatchNorm2d):
            nn.init.constant_(m.weight, 1)
            nn.init.constant_(m.bias, 0)

    # 保存权重
    weights_path = "model_weights.pt"
    torch.save(model.state_dict(), weights_path)

    # 验证文件大小
    file_size = os.path.getsize(weights_path)
    print(f"创建示例权重文件: {weights_path} (大小: {file_size / 1024:.1f}KB)")

    # 返回模型用于验证
    return model


if __name__ == "__main__":
    model = create_sample_weights()

    # 打印模型结构
    print("\n模型结构:")
    print(model)

    # 打印权重键
    print("\n权重键:")
    state_dict = model.state_dict()
    for key in state_dict.keys():
        print(f"- {key}: {state_dict[key].shape}")